package main

import (
	"container/heap"
	"fmt"
	"sort"
)

// 方法一   map+sort  O(nlogn)
func topKFrequent1(nums []int, k int) []int {
	mp := make(map[int]int)
	for _, num := range nums {
		mp[num]++
	}

	vec := make([]([2]int), 0)
	for k, v := range mp {
		vec = append(vec, [2]int{k, v})
	}
	sort.Slice(vec, func(i, j int) bool { return vec[i][1] > vec[j][1] })
	fmt.Println(vec)

	res := make([]int, 0)
	for i := 0; i < k; i++ {
		res = append(res, vec[i][0])
	}
	return res
}

type HeapVec [][2]int

func (vec HeapVec) Len() int {
	return len(vec)
}

func (vec HeapVec) Less(i, j int) bool {
	return vec[i][1] < vec[j][1]
}

func (vec HeapVec) Swap(i, j int) {
	vec[i], vec[j] = vec[j], vec[i]
}

func (vec *HeapVec) Push(x interface{}) {
	*vec = append(*vec, x.([2]int))
}

func (vec *HeapVec) Pop() interface{} {
	old := *vec
	n := len(old)
	item := old[n-1]
	*vec = old[0 : n-1]
	return item
}

// 方法二  优先级队列
func topKFrequent(nums []int, k int) []int {
	mp := make(map[int]int)
	for _, num := range nums {
		mp[num]++
	}
	fmt.Printf("mp: %v\n", mp)
	hp := &HeapVec{}
	heap.Init(hp)
	for key, value := range mp {
		heap.Push(hp, [2]int{key, value})
		if hp.Len() > k {
			heap.Pop(hp)
		}
		fmt.Println(hp)
	}

	res := make([]int, k)
	for i := 0; i < k; i++ {
		res[k-i-1] = heap.Pop(hp).([2]int)[0]
	}

	fmt.Println(res)
	return res

}

func main() {
	nums := []int{1, 1, 1, 2, 2, 3}
	topKFrequent(nums, 2)
}
